# @Author：zh
# @Data：2021/12/8 21:55
# @：PyCharm
# Python版本：3.7

## SoftMax 感知机分类的训练
## https://blog.csdn.net/weixin_43485035/article/details/115693919?spm=1001.2101.3001.6650.14&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-14.fixedcolumn&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-14.fixedcolumn
import numpy as np
import torch
from torchvision.datasets import mnist # 导入 pytorch 内置的 mnist 数据
from torch.utils.data import DataLoader
from torch import nn
from torch.autograd import Variable


## 全连接层
#对于神经网络，我们第一层的输入就是 28 x 28 = 784，所以必须将得到的数据我们做一个变换，使用 reshape 将他们拉平成一个一维向量
def data_tf(x):
    x = np.array(x, dtype='float32') / 255
    x = (x - 0.5) / 0.5 # 标准化，这个技巧之后会讲到
    x = x.reshape((-1,)) # 拉平成一个一维向量
    x = torch.from_numpy(x)
    return x

train_set = mnist.MNIST('\\minst', train=True, transform=data_tf, download=True) # 重新载入数据集，申明定义的数据变换
test_set = mnist.MNIST('\\minst', train=False, transform=data_tf, download=True)

## a表示单张图片  a_label表示图片的 标签 值
a, a_label = train_set[0]

# 使用 pytorch 自带的 DataLoader 定义一个数据迭代器
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
test_data = DataLoader(test_set, batch_size=128, shuffle=False)
#使用这样的数据迭代器是非常有必要的，如果数据量太大，就无法一次将他们全部读入内存，所以需要使用 python 迭代器，每次生成一个批次的数据

a, a_label = next(iter(train_data))

# 打印出一个批次的数据大小
print(a.shape)
print(a_label.shape)

## 定义网络的另一种方式 https://www.cnblogs.com/xuechengmeigui/p/12388400.html
# 使用 Sequential 定义 4 层神经网络(mlp)
net = nn.Sequential(
    ## 全连接层：输入 输出
    nn.Linear(784, 400),
    nn.ReLU(),
    nn.Linear(400, 200),
    nn.ReLU(),
    nn.Linear(200, 100),
    nn.ReLU(),
    nn.Linear(100, 10),##最后输出10个类别
    # nn.Softmax(dim=1)
)

# 定义 loss 函数：交叉熵
criterion = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(net.parameters(), 2e-1) # 使用随机梯度下降，学习率 0.1

# 开始训练
losses = []
acces = []
eval_losses = []
eval_acces = []

for e in range(20):
    train_loss = 0
    train_acc = 0
    net.train()
    for im, label in train_data:

        im = Variable(im)
        label = Variable(label)

        # 前向传播
        out = net(im)

        loss = criterion(out, label)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 记录误差
        train_loss += loss.item() ## loss值：https://www.zhihu.com/question/67209417/answer/344752405

        # 计算分类的准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / im.shape[0]
        train_acc += acc

    losses.append(train_loss / len(train_data))
    acces.append(train_acc / len(train_data))
    # 在测试集上检验效果
    eval_loss = 0
    eval_acc = 0
    net.eval() # 将模型改为预测模式
    for im, label in test_data:
        im = Variable(im)
        label = Variable(label)
        out = net(im)
        loss = criterion(out, label)
        # 记录误差
        eval_loss += loss.item()
        # 记录准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / im.shape[0]
        eval_acc += acc

    eval_losses.append(eval_loss / len(test_data))
    eval_acces.append(eval_acc / len(test_data))
    print('epoch: {}, Train Loss: {:.6f}, Train Acc: {:.6f}, Eval Loss: {:.6f}, Eval Acc: {:.6f}'.format(e, train_loss / len(train_data), train_acc / len(train_data),eval_loss / len(test_data), eval_acc / len(test_data)))

## 保存模型
torch.save(net.state_dict(), "Model_Minst_SoftMax.pt")

